#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI URL下载器插件
"""

import os
import sys
import time
import urllib.parse
import shutil
import requests
import re
from pathlib import Path
from typing import List, Tuple

# 导入统一的环境配置
from public.env_setup import setup_environment, ENV_CONFIG, URLDOWN_DIR
from public.download_utils import download_file_with_hf

# 使用统一的环境配置
ENV_CONFIG = setup_environment()

class URLDownloader:
    """URL下载器插件类"""
    
    # 插件排序
    order = 2
    
    def __init__(self):
        """初始化下载器"""
        # 延迟获取环境配置
        self.env_config = ENV_CONFIG
        self.workspace_path = self.env_config['PROJECT_ROOT']
        self.download_dir = URLDOWN_DIR
        self.max_retry = 3
        self.timeout = 30
        
        # 创建下载目录
        self.download_dir.mkdir(exist_ok=True)
        
        # 颜色定义
        self.RED = "\033[31m"
        self.GREEN = "\033[32m"
        self.YELLOW = "\033[33m"
        self.BLUE = "\033[34m"
        self.CYAN = "\033[36m"
        self.PURPLE = "\033[35m"
        self.RESET = "\033[0m"
    
    def start(self) -> bool:
        """启动下载器"""
        print(f"\n{self.PURPLE}■■■ URL下载器 v3.0（ComfyUI插件版） ■■■{self.RESET}")
        print(f"下载目录: {self.CYAN}{self.download_dir.absolute()}{self.RESET}")
        
        print("0. 退出")
        print("1) requests（推荐，总是可用）")
        print("2) curl（系统支持）")
        print("3) wget（不可用，需要安装）")
        tool = input("输入选项 (0/1/2/3): ").strip()
        if tool == "0" or tool.lower() == "q":
            print("已退出插件。")
            return False
        
        while True:
            print(f"\n{self.BLUE}════════════════════════════{self.RESET}")
            url = input("请输入下载URL (或输入 'q' 退出): ").strip()
            if url.lower() == "q":
                print(f"{self.GREEN}感谢使用，再见！{self.RESET}")
                break
            elif not url:
                print(f"{self.RED}错误：URL不能为空！{self.RESET}")
                continue

            filename = input("请输入保存文件名 (留空自动从URL提取): ").strip()
            self.download_file(url, filename, tool)
        
        return True
    
    @staticmethod
    def Launch(self):
        """静态启动方法"""
        downloader = URLDownloader()
        return downloader.start()
    
    def transform_huggingface_url(self, url: str) -> str:
        """转换HuggingFace URL为镜像地址"""
        if url.startswith("https://huggingface.co/") or url.startswith("https://hf-mirror.com/"):
            parsed = urllib.parse.urlparse(url)
            parts = parsed.path.strip("/").split("/")
            if len(parts) > 2:
                parts[2] = "resolve"

            # 替换 netloc 并重组路径
            new_path = "/" + "/".join(parts)
            new_parsed = parsed._replace(netloc="hf-mirror.com", path=new_path)
            return urllib.parse.urlunparse(new_parsed)
        return url

    def extract_filename_from_url(self, url: str) -> str:
        """从URL中提取文件名"""
        parsed_url = urllib.parse.urlparse(url)
        path = parsed_url.path
        basename = os.path.basename(path)
        filename = urllib.parse.unquote(basename)

        if '.' in filename and not filename.startswith('.'):
            return filename
        else:
            return f"{filename}.download"

    def check_tool_available(self, tool: str) -> bool:
        """检查下载工具是否可用"""
        try:
            if tool == "curl":
                result = os.system("curl --version >nul 2>&1")
                return result == 0
            elif tool == "wget":
                result = os.system("wget --version >nul 2>&1")
                return result == 0
            else:
                return True  # requests总是可用的
        except:
            return False

    def select_download_tool(self) -> str:
        """选择下载工具"""
        print(f"\n{self.CYAN}请选择下载工具：{self.RESET}")
        
        # 检查工具可用性
        curl_available = self.check_tool_available("curl")
        wget_available = self.check_tool_available("wget")
        
        print("1) requests（推荐，总是可用）")
        if curl_available:
            print("2) curl（系统支持）")
        else:
            print("2) curl（不可用）")
        if wget_available:
            print("3) wget（系统支持）")
        else:
            print("3) wget（不可用，需要安装）")
        
        choice = input("输入选项 (1/2/3): ").strip()
        
        if choice == "2" and curl_available:
            return "curl"
        elif choice == "3" and wget_available:
            return "wget"
        elif choice == "2" and not curl_available:
            print(f"{self.RED}curl不可用，使用requests{self.RESET}")
            return "requests"
        elif choice == "3" and not wget_available:
            print(f"{self.RED}wget不可用，使用requests{self.RESET}")
            return "requests"
        else:
            print(f"{self.YELLOW}默认使用 requests{self.RESET}")
            return "requests"

    def download_file(self, url: str, filename: str, tool: str):
        """下载文件"""
        url = self.transform_huggingface_url(url)
        auto_filename = self.extract_filename_from_url(url)

        if not filename:
            filename = auto_filename
            print(f"{self.YELLOW}未指定文件名，自动使用: {filename}{self.RESET}")
        elif '.' not in filename:
            ext = os.path.splitext(auto_filename)[1]
            filename += ext
            print(f"{self.YELLOW}已添加扩展名: {filename}{self.RESET}")

        filepath = self.download_dir / filename

        print(f"\n{self.BLUE}■ 下载信息 ■{self.RESET}")
        print(f"URL: {self.PURPLE}{url}{self.RESET}")
        print(f"保存位置: {self.PURPLE}{filepath}{self.RESET}")

        print(f"\n{self.CYAN}开始下载...{self.RESET}")
        success = False

        for attempt in range(1, self.max_retry + 1):
            print(f"{self.BLUE}▶ 尝试第 {attempt} 次 (最多 {self.max_retry} 次){self.RESET}")
            try:
                if tool == "curl":
                    print(f"{self.CYAN}使用curl下载...{self.RESET}")
                    result = os.system(f'curl -L -C - --connect-timeout {self.timeout} --progress-bar -o "{filepath}" "{url}"')
                    success = result == 0 and filepath.exists()
                elif tool == "wget":
                    print(f"{self.CYAN}使用wget下载...{self.RESET}")
                    result = os.system(f'wget -c --timeout={self.timeout} --progress=bar -O "{filepath}" "{url}"')
                    success = result == 0 and filepath.exists()
                else:
                    # 改进的requests下载
                    headers = {
                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                    }
                    
                    with requests.get(url, stream=True, timeout=self.timeout, headers=headers) as r:
                        r.raise_for_status()
                        total_size = int(r.headers.get('content-length', 0))
                        
                        print(f"{self.CYAN}文件大小: {total_size / (1024*1024):.2f} MB{self.RESET}")
                        
                        with open(filepath, 'wb') as f:
                            downloaded = 0
                            for chunk in r.iter_content(chunk_size=8192):
                                if chunk:
                                    f.write(chunk)
                                    downloaded += len(chunk)
                                    if total_size > 0:
                                        percent = (downloaded / total_size) * 100
                                        print(f"\r{self.CYAN}下载进度: {percent:.1f}% ({downloaded/(1024*1024):.2f}MB/{total_size/(1024*1024):.2f}MB){self.RESET}", end='', flush=True)
                        
                        print()  # 换行
                    success = True

                if success:
                    print(f"\n{self.CYAN}当前磁盘使用情况：{self.RESET}")
                    os.system("df -h")

                    break
            except Exception as e:
                print(f"{self.YELLOW}⚠ 下载失败：{e}，等待重试...{self.RESET}")
                time.sleep(3)

        if success:
            print(f"\n{self.GREEN}✓ 下载成功！{self.RESET}")
            try:
                size = filepath.stat().st_size
                print(f"{self.BLUE}文件大小: {size / 1024:.2f} KB{self.RESET}")
            except:
                pass
        else:
            print(f"\n{self.RED}✗ 下载失败！{self.RESET}")
            if filepath.exists():
                filepath.unlink()

def main():
    """主函数"""
    downloader = URLDownloader()
    downloader.start()

if __name__ == "__main__":
    main() 